#!/usr/bin/env python
import os
import sys
import shutil
import argparse
import random
import uuid
from lib.tables import message_table

config_dict={
	"kvm-img":"/usr/local/bin/kvm-img",
	"vir_dir":"/home/kong/VM",
	"mother":"/home/kong/VM/mother"
}

class Node_xml:
	def __init__(self):
		self.xml_example=os.path.split(os.path.realpath(__file__))[0]+"/template.xml"

	def make_uuid(self):
		return uuid.uuid4()

	def make_mac(self):
		mac = [ 0x52, 0x54, 0x00,random.randint(0x00, 0x7f),random.randint(0x00, 0xff),random.randint(0x00, 0xff) ]
		return ':'.join(map(lambda x: "%02x" % x, mac))

	def mkdir_xml(self,node,img,mem,cpu):
		if cpu == None:
			cpu=1

		try:
			c_size=int(cpu)
		except:
			print "cpu must be a int type!"
			return None

		if not 0 < c_size <=4:
			print "cpu size wrong!"
			return None

		if mem == None:
			mem="128M"

		try:
			m_size=int(mem[:-1])
		except:
			print "mem must be a int type!"
			return None

		if mem.endswith("G"):
			if not 0 < m_size <=4:
				print "mem size wrong!"
				return None
			m_size=m_size*1024*1024

		elif mem.endswith("M"):
			if not 128 <= m_size <=4096:
				print "mem size wrong!"
				return None
			m_size=m_size*1024

		else:
			print "mem size wrong!"
			return None


		mac=self.make_mac()
		uuid=self.make_uuid()
		xml=open(self.xml_example,'r').read()
		return xml % {"node":node,"uuid":uuid,"mem":m_size,"cpu":c_size,"img":img,"mac":mac}

class Kvm:
	def __init__(self):
		self.config=config_dict
		self.__check()

	def __check(self):
		if not self.config["vir_dir"].endswith("/"):
			self.config["vir_dir"]=self.config["vir_dir"]+"/"

		if not self.config["mother"].endswith("/"):
			self.config["mother"]=self.config["mother"]+"/"

		if not os.path.exists(self.config["vir_dir"]):
			print "kvm dir is not exists!"
			sys.exit(1)

		if not os.path.exists(self.config["mother"]):
			print "mother dir is not exists!"
			sys.exit(1)

		if not os.path.exists(self.config["kvm-img"]):
			print "kvm-img is not exists!"
			sys.exit(1)

	def list(self,**results):
		list_mother=results["list_mother"]

		if list_mother:
			self.__list_mother()
			return

		kvms=[(k+1,v) for k,v in enumerate(os.listdir(self.config["vir_dir"])) if v != "mother"]
		message_table(["id","node"],kvms)
		# for kvm in kvms:
		# 	print str(kvm[0])+" "+kvm[1]

	def __list_mother(self):
		imgs=[(k+1,v.split(".img")[0]) for k,v in enumerate(os.listdir(self.config["mother"]))]
		message_table(["id","img"],imgs)
		# for img in imgs:
		# 	print str(img[0])+" "+img[1]

	def reinit(self,node):
		if "," in node:
			nodes=node.split(",")
			for n in nodes:
				self.reinit(n)

		else:
			if not os.path.exists(self.config["vir_dir"]+node):
				print "node "+node+" is not exists!"
				sys.exit(1)
	
			node_path=self.config["vir_dir"]+node+"/"
			node_img=node_path+node+".img"
			node_img_bak=node_path+node+".img.bak"
			node_xml=node_path+node+".xml"
	
			if not (os.path.exists(node_img) and os.path.exists(node_img_bak) and os.path.exists(node_xml)):
				print node+" is not a node type!"
				sys.exit(1)
	
			os.remove(node_img)
			os.rename(node_img_bak, node_img)
			shutil.copy(node_img, node_img_bak)
			print node+" has been reinited!"

	def backup(self,node,force=False):
		node_path=self.config["vir_dir"]+node+"/"

		if not os.path.exists(node_path):
			print node+" is not existed!"
			sys.exit(1)

		node_img=node_path+node+".img"
		node_img_bak=node_path+node+".img.bak"

		if not os.path.exists(node_img):
			print node+" is not note type!"
			sys.exit(1)

		if not force:
			if os.path.exists(node_img_bak):
				print node+" has been backup before,try -f!"
				sys.exit(1)

		try:
			shutil.copy(node_img, node_img_bak)
			print node+" backup success!"
		except:
			print node+" backup failed!"
			sys.exit(1)


	def build(self,mother,node,mem,cpu):
		mother_path=self.config["mother"]+mother+".img"
		node_path=self.config["vir_dir"]+node+"/"

		if os.path.exists(node_path):
			print node+" has been existed!"
			sys.exit(1)

		if not os.path.exists(mother_path):
			print mother+" is not exists!"
			sys.exit(1)


		node_img=node_path+node+".img"
		#node_img_bak=node_path+node+".img.bak"
		node_xml=node_path+node+".xml"

		xml=Node_xml().mkdir_xml(node, node_img, mem, cpu)

		if not xml:
			sys.exit(1)

		os.mkdir(node_path)
		cmd="%s create -b %s -f qcow2 %s 2>&1 >> /dev/null" % (self.config["kvm-img"],mother_path,node_img)
		if os.system(cmd) != 0:
			print node_img+" create failed!"
			os.removedirs(node_path)
			sys.exit(1)

		#shutil.copy(node_img, node_img_bak)
		open(node_xml,'a').write(xml)
		print node_xml





def execute_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()

    p=subparsers.add_parser('list', help='list node')
    p.add_argument('--mother', action="store_true",dest="list_mother",default=False,help='list mother image')

    p=subparsers.add_parser('reinit', help='rebuild node')
    p.add_argument('node', action='store',help='node name')

    #mother,node,mem,cpu
    p=subparsers.add_parser('build', help='rebuild node')
    p.add_argument('mother', action='store',help='mother name')
    p.add_argument('node', action='store',help='node name')
    p.add_argument('-m', action='store',dest="mem",help='mem size')
    p.add_argument('-smp', action='store',dest="cpu",help='cpu size')

    p=subparsers.add_parser('backup', help='backup node')
    p.add_argument('node', action='store',help='node name')
    p.add_argument('-f','--force', action="store_true",dest="force",default=False,help='no asked')

    return parser


def execute_command(args=sys.argv[1:]):
    results = execute_args().parse_args(args)

    kvm=Kvm()

    commands = {
        'list': kvm.list,
        'reinit': kvm.reinit,
        'build': kvm.build,
        'backup': kvm.backup
    }

    func = commands[args[0]]
    func(**results.__dict__)

		


if __name__ == '__main__':
	execute_command()